//      motif_test_impl.cc
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#include <iostream>
#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp>
#include <boost/foreach.hpp>
#include <armadillo>

#include "stl-helper.h"
#include "../gene_processing.h"
#include "../gene_em.h"
#include "../gene_em_impl.h"

namespace po = boost::program_options;
extern po::variables_map vm;

void read_test()
{
  std::string input = vm["input"].as<std::string>();
  std::string annotationFile = vm["annotation"].as<std::string>();

  Genome genome;
  Annotation annotation;
  genome.read_genome(input);

  if (vm.count("simple"))
    annotation.read_annotation_simple(annotationFile);    
  else
    annotation.read_annotation(annotationFile);

  MSG("genome length = " << sz(genome));
  MSG("annotations = " << sz(annotation));

  annotation.get_gene_sequence(genome);
  if (vm.count("output")) { 
    std::string output = vm["output"].as<std::string>();
    annotation.save_upstream(output);
  }

  if (vm.count("pattern") && vm.count("freq")) {
    std::string freqFileName = vm["freq"].as<std::string>();
    std::string pattern = vm["pattern"].as<std::string>();
    
    Sequences upstreams;
    annotation.all_upstream(upstreams);
    
    std::vector<int> count;
    upstreams.count_occurence(pattern, count);
  
    std::ofstream freqFile;
    FileAccess::open2write(freqFile, freqFileName);
    
    freqFile << pattern;
    BOOST_FOREACH(int c, count) {
      freqFile << "\t " << c;
    }
    freqFile << ENDL;

    freqFile.close();
  }
}

/** Generate entropy of patterns
 *  input: a genome file and its annotation
 *  output: list of common patterns in the genome upstreams
 */
void entropy_test()
{
  std::string input = vm["input"].as<std::string>();

  Sequences upstreams;
  upstreams.read_sequence(input);

  MSG("sz(upstreams) = " << sz(upstreams));

  if (vm.count("output")) { 
    std::string output = vm["output"].as<std::string>();

    PatternGenerator pg(vm, upstreams, output);
    pg.run();
  }
}

/** Gibbs sampling for motif finding
 *  input: a genome file and its annotation
 *  output: motif in upstream sequences.
 *  freq: profile weight matrix
 */
void gibbs_test()
{
  std::string input = vm["input"].as<std::string>();

  Sequences upstreams;
  upstreams.read_sequence(input);

  MSG("n_seq = " << sz(upstreams));

  MotifGibbsSampler mgs(vm, upstreams);
  mgs.run();

  std::ofstream freq(vm["freq"].as<std::string>().c_str());
  const vvd_t& freq_ = mgs.freq();
  int motif_length = vm["mtlen"].as<int>();
  std::string acgt = "ACGT";
  std::string motif = std::string(motif_length, 'x');
  for (int c = 0; c < 4; c++) {
    freq << acgt[c];
    for (int pos = 0; pos <= motif_length; pos++) {
      bool ok = true;
      for (int c1 = 0; c1 < 4; c1++) 
	if (freq_[pos][c1] > freq_[pos][c]) ok = false;
      freq << "\t " << (ok?"+":"") << freq_[pos][c];
      if (ok && pos > 0) motif[pos-1] = acgt[c];
    }
    freq << ENDL;
  }
  freq << "SUM";
  for (int pos = 0; pos <= motif_length; pos++) {
    double sum = 0;
    for (int c = 0; c < 4; c++) 
      sum += freq_[pos][c];
    freq << "\t " << sum;
  }
  freq << ENDL << "motif = " << motif << ENDL
       << "likelihood = " << mgs.log_likelihood() << ENDL;
  freq.close();    
    
  std::ofstream output(vm["output"].as<std::string>().c_str());
  const vi_t& motif_position = mgs.pos();
  for (int i = 0; i < sz(upstreams); i++) {
    int pos = motif_position[i];
    output << ">" << upstreams.annotation_[i] << " motif=" << pos << ENDL;
    output << upstreams[i].substr(pos, motif_length) << ENDL;
  }
  output.close();

}

/** Motif statistics
 *  input/pattern: sequence file & motif 
 *  output: motif statistics 
 */
void motif_stats_test()
{
  std::string input = vm["input"].as<std::string>();
  std::string motif = vm["pattern"].as<std::string>();

  Sequences upstreams;
  upstreams.read_sequence(input);
  MSG("mt_stats n_seq = " << sz(upstreams) << " max-length = " << upstreams.max_length() 
      << " pattern = " << motif);

  MotifStats motif_stats(vm, upstreams);
  motif_stats.run();
  motif_stats.output();
}

void em_test()
{
  std::string input = vm["input"].as<std::string>();
  Sequences seqs;

  seqs.read_sequence(input);
  MSG("EM n_seq = " << sz(seqs) << " max-length = " << seqs.max_length()); 

  MotifModel max_model;
  double max = -DOUBLE_INFINITY;
  Accumulate<MotifModel, std::less<double> > acc(max, max_model);
  
  for (int i = 0; i < vm["run"].as<int>(); i++) {
    MotifModel model;
    int motif_length = vm["mtlen"].as<int>();
    model.random_initialize(motif_length, seqs.max_length());

    int max_iter = vm["iter"].as<int>();
    double likelihood = gene_em<MotifModel, Sequences>(max_iter, model, seqs);
    
    acc(likelihood, model);
  }
  
  std::string output = vm["output"].as<std::string>();
  max_model.param().save(output);
  MSG("likelihood = " << max);
}

void em_max_test()
{
  std::string input = vm["input"].as<std::string>();
  Sequences seqs;
  GroupSequences group_seqs, group_seqs1;

  seqs.read_sequence(input);
  //BOOST_FOREACH(std::string s, seqs) MSG(s);
  
  group_seqs1.group_sequences(seqs);
  group_seqs = group_seqs1.removeSharp();
  MSG("EM n_seq = " << sz(seqs) << " max-length = " << seqs.max_length()); 
  MSG("n_group = " << sz(group_seqs));

  MotifModel max_model;
  double max = -DOUBLE_INFINITY;
  Accumulate<MotifModel> acc(max, max_model);
  for (int i = 0; i < vm["run"].as<int>(); i++) {
    MSG("============= " << i << " ===============");
    MotifModel model;
    int motif_length = vm["mtlen"].as<int>();
    model.random_initialize(motif_length, seqs.max_length());

    int max_iter = vm["iter"].as<int>();
    double likelihood = gene_em_max<MotifModel, GroupSequences>(max_iter, model, group_seqs);
    
    acc(likelihood, model);
  }  
  std::string output = vm["output"].as<std::string>();
  max_model.param().save(output);
  
}

void extract_upstream()
{
  std::string geneFile = vm["input"].as<std::string>();
  std::string annoFile = vm["annotation"].as<std::string>();
  std::string outputFile = vm["output"].as<std::string>();
  
  Genome genome;
  Annotation annotation;
  
  genome.read_genome(geneFile);
  annotation.read_annotation_simple1(annoFile);
  
  annotation.get_gene_sequence(genome, vm["seqlen"].as<int>());
  annotation.save_upstream(outputFile);
}

void new_motif_stats_test()
{
  std::string seqFile = vm["input"].as<std::string>();
  std::string outputFile = vm["output"].as<std::string>();
  int mtlen = vm["mtlen"].as<int>();
  
  Sequences seqs;
  seqs.read_sequence(seqFile);
  int maxlen = seqs.max_length();
  
  std::map<std::string, vi_t> mt_freq;
  typedef std::map<std::string, vi_t>::value_type value_type;
  
  struct add_motif {
    std::map<std::string, vi_t>& mt_freq;
    int maxlen;
    
    add_motif(std::map<std::string, vi_t>& mt_freq_, int maxlen_) 
      : mt_freq(mt_freq_), maxlen(maxlen_) {}
    
    void operator() (const std::string& motif, int pos) {
      if (mt_freq.find(motif) == mt_freq.end()) 
          mt_freq[motif] = vi_t(maxlen, 0);
        mt_freq[motif][pos] ++;
    }
  };
  
  add_motif add(mt_freq, maxlen);
  BOOST_FOREACH(const std::string& seq, seqs) {
    int len = sz(seq);
    for (int i = 0; i < len; i++) {
      for (int j = 1; j < mtlen; j++) if (i + j <= len) {
        std::string motif = seq.substr(i, j);
        add(motif, i);
        for (int k = 0; k < sz(motif); k++) {
          char tmp = motif[k];
          motif[k] = 'x';
          add(motif, i);
          motif[k] = tmp;
        }
      }
    }
  }
  
  std::ofstream out;
  FileAccess::open2write(out, outputFile);
  
  out << "motif";
  for (int i = 0; i < maxlen; i++) out << "\t pos(" << i-maxlen << ")";
  out << ENDL;
  
  BOOST_FOREACH(const value_type& freq, mt_freq) {
    out << freq.first;
    BOOST_FOREACH(int count, freq.second) out << "\t " << count;
    out << ENDL;
  }
  out.close();
}
/*
void em2_test()
{
  std::string input = vm["input"].as<std::string>();
  Sequences seqs;

  seqs.read_sequence(input);
  MSG("EM n_seq = " << sz(seqs) << " max-length = " << seqs.max_length()); 

  Motif2ndModel max_model;
  double max = -DOUBLE_INFINITY;
  Accumulate<Motif2ndModel, std::less<double> > acc(max, max_model);
  
  for (int i = 0; i < vm["run"].as<int>(); i++) {
    Motif2ndModel model;
    int motif_length = vm["mtlen"].as<int>();
    model.random_initialize(motif_length, seqs.max_length());

    int max_iter = vm["iter"].as<int>();
    double likelihood = gene_em<Motif2ndModel, Sequences>(max_iter, model, seqs);
    
    acc(likelihood, model);
  }
  
  std::string output = vm["output"].as<std::string>();
  max_model.param().save(output);
  MSG("likelihood = " << max);
}

*/
